"""
Convert Three.js HTML animation to MP4 video.
"""

import logging
import time
from pathlib import Path
from typing import Tuple, Dict, Any
from selenium import webdriver
from selenium.webdriver.chrome.service import Service
from selenium.webdriver.chrome.options import Options
from webdriver_manager.chrome import ChromeDriverManager
import cv2
import numpy as np
from tqdm import tqdm

def capture_frames(
    html_file: str,
    output_dir: str,
    frames: int,
    resolution: Tuple[int, int],
    config: Dict[str, Any]
) -> str:
    """
    Capture frames from Three.js animation and save them as MP4.
    
    Args:
        html_file: Path to HTML file
        output_dir: Directory to save the MP4 file
        frames: Number of frames to capture
        resolution: Video resolution (width, height)
        config: Configuration dictionary
    
    Returns:
        Path to the generated MP4 file
    """
    driver = None
    try:
        logging.info("Initializing Chrome WebDriver...")
        options = webdriver.ChromeOptions()
        options.add_argument("--headless=new")
        options.add_argument("--no-sandbox")
        options.add_argument("--disable-dev-shm-usage")
        options.add_argument(f"--window-size={resolution[0]},{resolution[1]}")
        
        service = Service(ChromeDriverManager().install())
        driver = webdriver.Chrome(service=service, options=options)
        
        # 加载HTML文件
        file_url = f"file://{Path(html_file).absolute()}"
        logging.info(f"Loading file: {file_url}")
        driver.get(file_url)
        time.sleep(2)  # 等待页面加载
        
        # 准备视频写入器
        output_path = Path(output_dir) / "output.mp4"
        fourcc = cv2.VideoWriter_fourcc(*'mp4v')
        out = cv2.VideoWriter(
            str(output_path),
            fourcc,
            30,  # FPS
            resolution
        )
        
        logging.info("Starting frame capture...")
        with tqdm(total=frames, desc="Capturing frames") as pbar:
            for frame in range(frames):
                # 更新动画状态（如果需要）
                driver.execute_script(f"animate({frame});")
                
                # 捕获截图
                screenshot = driver.get_screenshot_as_png()
                nparr = np.frombuffer(screenshot, np.uint8)
                img = cv2.imdecode(nparr, cv2.IMREAD_COLOR)
                
                # 调整图像大小
                img = cv2.resize(img, resolution)
                
                # 写入帧
                out.write(img)
                pbar.update(1)
        
        out.release()
        logging.info(f"Video saved to: {output_path}")
        return str(output_path)
        
    except Exception as e:
        logging.error(f"Error during video capture: {str(e)}", exc_info=True)
        raise
    finally:
        if driver:
            driver.quit()

def main():
    """Main function"""
    try:
        from utils import setup_logging, load_config, get_project_root
        
        # 加载配置
        config_path = get_project_root() / "config" / "settings.json"
        config = load_config(config_path)
        
        # 设置日志
        setup_logging(config)
        
        logging.info("=== Starting HTML to MP4 conversion ===")
        
        # 确保输出目录存在
        output_dir = Path("output")
        output_dir.mkdir(exist_ok=True)
        
        # 捕获帧并生成视频
        video_path = capture_frames(
            html_file=config['html_file'],
            output_dir=str(output_dir),
            frames=config['frames'],
            resolution=tuple(config['resolution']),
            config=config
        )
        
        logging.info("=== Conversion completed successfully! ===")
        logging.info(f"Output video: {video_path}")
        
    except Exception as e:
        logging.error(f"Conversion failed: {str(e)}", exc_info=True)
        raise

if __name__ == "__main__":
    main() 